from collections import deque
import xxhash
import numpy as np

from nanovllm.engine.sequence import Sequence


class Block:

    def __init__(self, block_id):
        self.block_id = block_id # 块的唯一ID
        self.ref_count = 0 # 引用计数（有多少序列在使用该块）
        self.hash = -1 # 块内容的哈希值（用于复用）
        self.token_ids = []  # 该块存储的token列表

    def update(self, hash: int, token_ids: list[int]):
        """更新块的内容（哈希和token列表）"""
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        """重置块（释放时用，清空状态）"""
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []


class BlockManager:
    """
    核心作用：封装单个 KV 缓存块的状态（是否被使用、存储的内容），支持复用和释放。
    ref_count：多个序列可能共享同一块（如不同序列的提示词有重叠时），引用计数确保块在被完全释放前不被销毁。
    """

    def __init__(self, num_blocks: int, block_size: int):
        assert num_blocks > 0
        self.block_size = block_size # 每个块的大小（如256 token）
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)] # 初始化所有块
        self.hash_to_block_id: dict[int, int] = dict() # 哈希到块ID的映射（用于快速查找复用块）
        self.free_block_ids: deque[int] = deque(range(num_blocks)) # 空闲块ID队列（快速分配）
        self.used_block_ids: set[int] = set()  # 已使用块ID集合（快速判断状态）

    @classmethod
    def compute_hash(cls, token_ids: list[int], prefix: int = -1):
        """
        作用：为一组 token 生成唯一哈希，用于判断是否已有相同内容的缓存块（实现 “缓存复用”）
        前缀哈希：若块内容依赖前一个块（如生成过程中，当前块的 token 依赖前一个块的结果），则加入前缀哈希，确保链式依赖的正确性。
        :param token_ids:
        :param prefix:
        :return: intdigest
        """
        h = xxhash.xxh64() # 使用xxhash快速计算哈希
        if prefix != -1:
            h.update(prefix.to_bytes(8, "little"))  # 前缀哈希（用于链式依赖）
        h.update(np.array(token_ids).tobytes()) # 将token列表转为字节流计算哈希
        return h.intdigest()  # 返回整数哈希值

    def _allocate_block(self, block_id: int) -> Block:
        block = self.blocks[block_id]
        assert block.ref_count == 0 # 确保块当前未被使用
        block.reset() # 重置块状态
        self.free_block_ids.remove(block_id) # 从空闲队列移除
        self.used_block_ids.add(block_id) # 标记为已使用
        return self.blocks[block_id]

    def _deallocate_block(self, block_id: int) -> Block:
        assert self.blocks[block_id].ref_count == 0 # 确保引用计数为0（无序列使用）
        self.used_block_ids.remove(block_id) # 标记为未使用
        self.free_block_ids.append(block_id) # 加入空闲队列

    def can_allocate(self, seq: Sequence) -> bool:
        # 判断空闲块数量是否 >= 序列需要的块数量
        return len(self.free_block_ids) >= seq.num_blocks

    def allocate(self, seq: Sequence):
        """
        缓存复用：通过哈希查找是否已有相同内容的块，若存在则复用（引用计数 + 1），否则分配新块。
        链式哈希：通过 prefix=h 传递前一个块的哈希，确保连续块的依赖关系（如生成过程中，当前块的内容依赖前一个块）。
        seq.block_table：记录序列使用的所有块 ID，后续释放时需遍历该列表。
        :param seq:
        :return:
        """
        assert not seq.block_table # 确保序列未分配过块
        h = -1
        cache_miss = False
        for i in range(seq.num_blocks):
            token_ids = seq.block(i) # 获取序列的第i个块的token列表
            # 计算哈希（若token数等于块大小，说明是完整块，加入前缀哈希）
            h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
            block_id = self.hash_to_block_id.get(h, -1) # 查找是否已有相同哈希的块
            # 若块不存在或内容不匹配，则标记为缓存未命中（需要分配新块）
            if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
                cache_miss = True
            # 缓存未命中：分配新块
            if cache_miss:
                block_id = self.free_block_ids[0] # 取第一个空闲块
                block = self._allocate_block(block_id) # 分配块
            else:
                # 缓存命中：复用已有块，增加引用计数
                seq.num_cached_tokens += self.block_size # 缓存token数+块大小
                if block_id in self.used_block_ids:
                    block = self.blocks[block_id]
                    block.ref_count += 1  # 引用计数+1
                else:
                    block = self._allocate_block(block_id)  # 若块未被使用，直接分配
            if h != -1: # 若哈希有效，更新块的哈希和token列表，并记录到哈希映射
                block.update(h, token_ids)
                self.hash_to_block_id[h] = block_id
            seq.block_table.append(block_id) # 将块ID加入序列的block_table（记录该序列使用的块）

    def deallocate(self, seq: Sequence):
        """
        反向遍历：确保先释放后续的块（可能依赖前面的块），避免提前释放被依赖的块。
        引用计数管理：只有当所有使用该块的序列都释放后（ref_count == 0），才真正将块加入空闲队列，支持 “多序列复用同一块” 的场景。
        :param seq:
        :return:
        """
        for block_id in reversed(seq.block_table): # 反向遍历序列的block_table（确保依赖顺序正确释放）
            block = self.blocks[block_id]
            block.ref_count -= 1 # 引用计数-1
            if block.ref_count == 0: # 若引用计数为0，释放该块（加入空闲队列）
                self._deallocate_block(block_id)
        seq.num_cached_tokens = 0  # 重置序列的缓存状态
        seq.block_table.clear() # 清空block_table

    def can_append(self, seq: Sequence) -> bool:
        # 判断是否需要新块：若当前token数是块大小的整数倍+1（如257=256+1），则需要新块
        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

    def may_append(self, seq: Sequence):
        """
        若需要新块（len(seq) % block_size == 1）：分配新块并加入 block_table。
        若最后一个块已填满（len(seq) % block_size == 0）：计算该块的哈希并更新内容，标记为可用。
        :param seq:
        :return:
        """
        block_table = seq.block_table
        last_block = self.blocks[block_table[-1]] # 获取序列的最后一个块
        if len(seq) % self.block_size == 1:  # 若当前token数是块大小+1（需要新块）
            assert last_block.hash != -1 # 确保最后一个块已初始化
            block_id = self.free_block_ids[0]  # 取空闲块
            self._allocate_block(block_id) # 分配新块
            block_table.append(block_id) # 加入block_table
        elif len(seq) % self.block_size == 0: # 若当前token数是块大小的整数倍（最后一个块已填满）
            assert last_block.hash == -1 # 确保最后一个块未初始化
            token_ids = seq.block(seq.num_blocks-1) # 获取最后一个块的token列表
            prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1 # 计算前缀哈希（若有前一个块，用其哈希；否则为-1）
            h = self.compute_hash(token_ids, prefix)  # 计算哈希
            last_block.update(h, token_ids) # 更新最后一个块的内容
            self.hash_to_block_id[h] = last_block.block_id # 记录哈希映射
        else:
            assert last_block.hash == -1 # 其他情况，确保最后一个块未初始化


    """
    BlockManager 是 LLM 推理中KV 缓存复用的核心，通过以下设计优化性能：
    
    缓存复用：用哈希快速查找重复块，避免重复计算 KV 向量，大幅减少计算量。
    引用计数：支持多序列复用同一块（如不同对话的提示词有重叠时），提升缓存利用率。
    动态块管理：按需分配 / 释放块，结合 Scheduler 实现高并发任务的资源调度。
    
    理解 BlockManager 后，就能明白 LLM 推理中 “为什么长文本生成更快”（因为后续 token 复用了前面的 KV 缓存块），以及 “如何支持高并发”（通过块复用和引用计数优化资源）。
    """
